from fastapi import APIRouter, HTTPException, UploadFile, File, Form, Query, BackgroundTasks, status, Request
from fastapi.responses import FileResponse, JSONResponse
import os
import shutil
from typing import Optional, Dict, Any
from pathlib import Path
import uuid
from pydantic import BaseModel
import asyncio
import threading
import time

from .translateFile import translate_pdf_file, translate_docx_file, translate_epub_file

router = APIRouter(tags=["文件翻译"])

class FileTranslationResponse(BaseModel):
    file_url: str  # EPUB/PDF/DOCX文件URL
    message: str
    filename: str
    output_path: str
    total_pages: Optional[int] = None
    translated_pages: Optional[int] = None
    from_page: Optional[int] = None
    to_page: Optional[int] = None
    job_id: Optional[str] = None  # 添加作业ID，用于异步查询进度
    
    class Config:
        schema_extra = {
            "example": {
                "file_url": "http://example.com/translate/file/download/document_zh-CN.epub",
                "message": "文件翻译成功",
                "filename": "document.epub",
                "output_path": "output_files/document_zh-CN.epub",
                "job_id": "550e8400-e29b-41d4-a716-446655440000"
            }
        }

class ProgressResponse(BaseModel):
    job_id: str
    status: str  # "pending", "processing", "completed", "failed"
    progress: int  # 0-100
    message: Optional[str] = None
    file_url: Optional[str] = None

# 存储翻译作业的状态
translation_jobs: Dict[str, Dict[str, Any]] = {}

@router.post("/pdf", response_model=FileTranslationResponse)
async def translate_pdf(
    request: Request,  # 添加请求对象以获取基础URL
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    source_lang: str = Form("auto"),
    target_lang: str = Form("zh-CN"),
    from_page: int = Form(1, description="起始页码（从1开始）"),
    to_page: Optional[int] = Form(None, description="结束页码（包含该页，默认为最后一页）"),
    simple_response: bool = Form(False, description="如果为True，仅返回文件URL"),
    only_txt: bool = Form(False, description="如果为True，仅返回文本文件URL")
):
    """
    翻译PDF文件（支持多页，可指定页码范围，保持PDF格式）
    
    - **file**: 要翻译的PDF文件
    - **source_lang**: 源语言代码，默认为auto自动检测
    - **target_lang**: 目标语言代码，默认为zh-CN(简体中文)
    - **from_page**: 起始页码（从1开始），默认为1
    - **to_page**: 结束页码（包含该页），默认为None（表示最后一页）
    - **simple_response**: 设置为True时仅返回文件URL，默认为False
    - **only_txt**: 设置为True时仅返回文本文件URL，默认为False
    
    注意：翻译大文件可能需要较长时间，请耐心等待。
    返回结果包含PDF文件URL和文本文件URL。
    """
    try:
        # 检查文件是否为PDF
        if not file.filename.lower().endswith('.pdf'):
            raise HTTPException(status_code=400, detail="仅支持PDF文件")
        
        # 验证页码范围
        if from_page < 1:
            raise HTTPException(status_code=400, detail=f"起始页码必须大于等于1，当前值: {from_page}")
        
        if to_page is not None and to_page < from_page:
            raise HTTPException(status_code=400, detail=f"结束页码必须大于等于起始页码，当前起始页码: {from_page}，结束页码: {to_page}")
        
        # 创建唯一文件名
        unique_id = str(uuid.uuid4())
        input_folder = Path("input_files")
        input_folder.mkdir(exist_ok=True)
        
        # 保存上传的文件
        file_path = input_folder / f"{unique_id}_{file.filename}"
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        # 创建作业ID
        job_id = str(uuid.uuid4())
        
        # 初始化作业状态
        translation_jobs[job_id] = {
            "status": "processing",
            "file_path": str(file_path),
            "filename": file.filename,
            "progress": 0
        }
        
        # 创建输出路径，确保是PDF格式
        output_folder = Path("output_files")
        output_folder.mkdir(exist_ok=True)
        output_filename = f"{Path(file.filename).stem}_{target_lang}.pdf"
        output_path = output_folder / output_filename
        
        # 翻译文件
        translated_text, saved_path, txt_output_path = translate_pdf_file(
            input_path=str(file_path),
            output_path=str(output_path),
            source_lang=source_lang,
            target_lang=target_lang,
            from_page=from_page,
            to_page=to_page
        )
        
        # 更新作业状态
        translation_jobs[job_id]["status"] = "completed"
        translation_jobs[job_id]["progress"] = 100
        translation_jobs[job_id]["output_path"] = saved_path
        
        # 在完成后删除临时文件
        def cleanup_temp_file():
            if os.path.exists(file_path):
                os.remove(file_path)
            # 60分钟后删除作业状态
            # 在实际应用中，你可能需要一个定时任务来清理旧的作业
        
        background_tasks.add_task(cleanup_temp_file)
        
        # 计算翻译页数
        actual_to_page = to_page
        actual_from_page = from_page
        
        # 提取文件名并创建URL
        filename_only = os.path.basename(saved_path)
        txt_filename_only = os.path.basename(txt_output_path)
        base_url = str(request.base_url).rstrip('/')
        file_url = f"{base_url}/translate/file/download/{filename_only}"
        txt_file_url = f"{base_url}/translate/file/download/{txt_filename_only}"
        
        # 检查是否只返回TXT URL
        if only_txt:
            return JSONResponse(content={"txt_file_url": txt_file_url})
            
        # 检查是否只返回URL
        if simple_response:
            return JSONResponse(content={"file_url": file_url, "txt_file_url": txt_file_url})
        
        # 返回完整响应
        return FileTranslationResponse(
            file_url=file_url,
            message=f"PDF文件翻译成功，已处理页码范围 {from_page} - {actual_to_page if actual_to_page else '末页'}",
            filename=file.filename,
            output_path=saved_path,
            from_page=from_page,
            to_page=actual_to_page
        )
    
    except ValueError as e:
        # 处理值错误（例如页码范围错误）
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        # 异常处理
        raise HTTPException(status_code=500, detail=f"PDF文件翻译失败: {str(e)}")

@router.post("/docx", response_model=FileTranslationResponse)
async def translate_docx(
    request: Request,  # 添加请求对象以获取基础URL
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    source_lang: str = Form("auto"),
    target_lang: str = Form("zh-CN"),
    simple_response: bool = Form(False, description="如果为True，仅返回文件URL"),
    only_txt: bool = Form(False, description="如果为True，仅返回文本文件URL")
):
    """
    翻译DOCX文件（保持DOCX格式）
    
    - **file**: 要翻译的DOCX文件
    - **source_lang**: 源语言代码，默认为auto自动检测
    - **target_lang**: 目标语言代码，默认为zh-CN(简体中文)
    - **simple_response**: 设置为True时仅返回文件URL，默认为False
    - **only_txt**: 设置为True时仅返回文本文件URL，默认为False
    
    注意：翻译大文件可能需要较长时间，请耐心等待。
    返回结果包含DOCX文件URL和文本文件URL。
    """
    try:
        # 检查文件是否为DOCX
        if not file.filename.lower().endswith('.docx'):
            raise HTTPException(status_code=400, detail="仅支持DOCX文件")
        
        # 创建唯一文件名
        unique_id = str(uuid.uuid4())
        input_folder = Path("input_files")
        input_folder.mkdir(exist_ok=True)
        
        # 保存上传的文件
        file_path = input_folder / f"{unique_id}_{file.filename}"
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        # 设置输出路径，确保是DOCX格式
        output_folder = Path("output_files")
        output_folder.mkdir(exist_ok=True)
        output_filename = f"{Path(file.filename).stem}_{target_lang}.docx"
        output_path = output_folder / output_filename
        
        # 翻译文件
        translated_text, saved_path, txt_output_path = translate_docx_file(
            input_path=str(file_path),
            output_path=str(output_path),
            source_lang=source_lang,
            target_lang=target_lang
        )
        
        # 在完成后删除临时文件
        def cleanup_temp_file():
            if os.path.exists(file_path):
                os.remove(file_path)
        
        background_tasks.add_task(cleanup_temp_file)
        
        # 提取文件名并创建URL
        filename_only = os.path.basename(saved_path)
        txt_filename_only = os.path.basename(txt_output_path)
        base_url = str(request.base_url).rstrip('/')
        file_url = f"{base_url}/translate/file/download/{filename_only}"
        txt_file_url = f"{base_url}/translate/file/download/{txt_filename_only}"
        
        # 检查是否只返回TXT URL
        if only_txt:
            return JSONResponse(content={"txt_file_url": txt_file_url})
            
        # 检查是否只返回URL
        if simple_response:
            return JSONResponse(content={"file_url": file_url, "txt_file_url": txt_file_url})
            
        # 返回完整响应
        return FileTranslationResponse(
            file_url=file_url,
            message="DOCX文件翻译成功",
            filename=file.filename,
            output_path=saved_path,
            total_pages=None,
            translated_pages=None,
            from_page=None,
            to_page=None
        )
    
    except Exception as e:
        # 异常处理
        raise HTTPException(status_code=500, detail=f"DOCX文件翻译失败: {str(e)}")

@router.post("/epub", response_model=FileTranslationResponse)
async def translate_epub(
    request: Request,
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    source_lang: str = Form("auto"),
    target_lang: str = Form("zh-CN"),
    simple_response: bool = Form(False, description="如果为True，仅返回文件URL"),
    only_txt: bool = Form(False, description="如果为True，仅返回文本文件URL"),
    translate_content: bool = Form(True, description="如果为True，翻译EPUB内容")
):
    """
    翻译EPUB电子书文件（翻译后保持EPUB格式）
    
    - **file**: 要翻译的EPUB文件
    - **source_lang**: 源语言代码，默认为auto自动检测
    - **target_lang**: 目标语言代码，默认为zh-CN(简体中文)
    - **simple_response**: 设置为True时仅返回文件URL，默认为False
    - **only_txt**: 设置为True时仅返回文本文件URL，默认为False
    - **translate_content**: 如果为True，修改EPUB文件内容并翻译，默认为True
    
    注意：翻译大文件可能需要较长时间，请耐心等待。
    翻译将在后台进行，您可以通过返回的job_id查询进度。
    结果将包含翻译后的EPUB文件URL。
    """
    try:
        # 检查文件是否为EPUB
        if not file.filename.lower().endswith('.epub'):
            raise HTTPException(status_code=400, detail="仅支持EPUB文件")
        
        # 创建唯一文件名
        unique_id = str(uuid.uuid4())
        input_folder = Path("input_files")
        input_folder.mkdir(exist_ok=True)
        
        # 保存上传的文件
        file_path = input_folder / f"{unique_id}_{file.filename}"
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        # 创建作业ID
        job_id = str(uuid.uuid4())
        
        # 设置输出路径，确保是EPUB格式
        output_folder = Path("output_files")
        output_folder.mkdir(exist_ok=True)
        output_filename = f"{Path(file.filename).stem}_{target_lang}.epub"
        output_path = output_folder / output_filename
        txt_output_path = str(output_path).replace('.epub', '.txt')
        
        # 初始化作业状态
        translation_jobs[job_id] = {
            "status": "pending",
            "file_path": str(file_path),
            "filename": file.filename,
            "progress": 0,
            "output_path": str(output_path),
            "txt_output_path": None,  # 设置为None，因为我们不会保留txt文件
            "message": "翻译任务已排队"
        }
        
        # 计算URL
        base_url = str(request.base_url).rstrip('/')
        file_url = f"{base_url}/translate/file/download/{output_filename}"
        
        # 定义进度回调函数
        def update_progress(progress):
            if job_id in translation_jobs:
                translation_jobs[job_id]["progress"] = progress
                if progress == -1:  # 翻译失败
                    translation_jobs[job_id]["status"] = "failed"
                    translation_jobs[job_id]["message"] = "翻译失败"
                elif progress == 100:  # 翻译完成
                    translation_jobs[job_id]["status"] = "completed"
                    translation_jobs[job_id]["message"] = "翻译完成"
                else:  # 翻译进行中
                    translation_jobs[job_id]["status"] = "processing"
                    translation_jobs[job_id]["message"] = f"正在翻译 ({progress}%)"
        
        # 定义后台翻译任务
        def background_translate():
            try:
                translation_jobs[job_id]["status"] = "processing"
                translation_jobs[job_id]["message"] = "正在翻译"
                
                # 调用翻译函数
                translated_text, saved_path, txt_path = translate_epub_file(
                    input_path=str(file_path),
                    output_path=str(output_path),
                    source_lang=source_lang,
                    target_lang=target_lang,
                    progress_callback=update_progress
                )
                
                # 更新作业状态
                if job_id in translation_jobs:
                    translation_jobs[job_id]["status"] = "completed"
                    translation_jobs[job_id]["progress"] = 100
                    translation_jobs[job_id]["output_path"] = saved_path
                    translation_jobs[job_id]["message"] = "翻译完成"
                    
                # 在完成后删除临时文件
                if os.path.exists(file_path):
                    os.remove(file_path)
                    
                # 删除txt文件，只保留epub
                if os.path.exists(txt_path):
                    os.remove(txt_path)
                    
            except Exception as e:
                # 异常处理
                if job_id in translation_jobs:
                    translation_jobs[job_id]["status"] = "failed"
                    translation_jobs[job_id]["progress"] = -1
                    translation_jobs[job_id]["message"] = f"翻译失败: {str(e)}"
                print(f"EPUB文件翻译失败: {str(e)}")
                
                # 确保临时文件被删除
                if os.path.exists(file_path):
                    os.remove(file_path)
        
        # 启动后台翻译任务
        # 使用线程而不是直接用background_tasks，以便更好地控制任务和提供进度更新
        threading.Thread(target=background_translate).start()
        
        # 设置作业清理的后台任务（仅清理作业状态，不涉及实际翻译）
        def cleanup_job():
            # 24小时后从字典中删除作业状态
            time.sleep(24 * 60 * 60)
            if job_id in translation_jobs:
                del translation_jobs[job_id]
        
        background_tasks.add_task(cleanup_job)
        
        # 简化响应，仅返回EPUB文件URL和作业ID
        return_data = {
            "file_url": file_url,
            "job_id": job_id,
            "message": "EPUB翻译任务已开始，将在后台进行"
        }
        
        # 检查是否只返回简单响应或URL
        if simple_response or only_txt:
            return JSONResponse(content=return_data)
            
        # 返回完整响应
        return FileTranslationResponse(
            file_url=file_url,
            message="EPUB文件翻译已开始，将在后台进行",
            filename=file.filename,
            output_path=str(output_path),
            job_id=job_id,
            total_pages=None,
            translated_pages=None,
            from_page=None,
            to_page=None
        )
    
    except Exception as e:
        # 异常处理
        raise HTTPException(status_code=500, detail=f"EPUB文件翻译失败: {str(e)}")

@router.get("/progress/{job_id}", response_model=ProgressResponse)
async def get_translation_progress(
    request: Request,
    job_id: str
):
    """
    获取翻译作业的进度
    
    - **job_id**: 作业ID，由翻译API返回
    
    返回翻译作业的当前状态、进度百分比和消息。
    如果作业已完成，还将包含翻译后的EPUB文件URL。
    """
    if job_id not in translation_jobs:
        raise HTTPException(status_code=404, detail=f"找不到作业: {job_id}")
    
    job = translation_jobs[job_id]
    
    # 计算URL（如果作业已完成）
    file_url = None
    
    if job["status"] in ["completed", "failed"] and "output_path" in job:
        base_url = str(request.base_url).rstrip('/')
        filename_only = os.path.basename(job["output_path"])
        file_url = f"{base_url}/translate/file/download/{filename_only}"
    
    return ProgressResponse(
        job_id=job_id,
        status=job["status"],
        progress=job["progress"],
        message=job.get("message", ""),
        file_url=file_url
    )

@router.get("/download/{filename:path}")
async def download_translated_file(filename: str):
    """
    下载翻译后的文件
    
    - **filename**: 文件名称
    """
    file_path = Path("output_files") / filename
    
    if not file_path.exists():
        raise HTTPException(status_code=404, detail=f"文件不存在: {filename}")
    
    # 根据文件扩展名设置正确的MIME类型
    media_type = "text/plain"  # 默认为文本
    if filename.lower().endswith('.pdf'):
        media_type = "application/pdf"
    elif filename.lower().endswith('.docx'):
        media_type = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
    elif filename.lower().endswith('.epub'):
        media_type = "application/epub+zip"
    
    return FileResponse(
        path=str(file_path),
        filename=filename,
        media_type=media_type
    ) 